package net.cassandraview.db;

import org.apache.cassandra.thrift.*;
import org.apache.thrift.TException;

import java.nio.ByteBuffer;
import java.nio.charset.Charset;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.logging.Level;
import java.util.logging.Logger;

/**
 * Connection to Cassandra database using thrift api.
 *
 * User: nagakhl
 * Date: 21.01.12
 * Time: 21:38
 */
public class ThriftConnection implements Connection {
    
    protected Logger log = Logger.getLogger(getClass().getName());
    
    protected String keyspace;
    
    protected int defaultLimit;

    protected ConsistencyLevel defaultLevel;

    protected Cassandra.Client client;

    protected ConnectionManager parentManager;

    protected Throwable lastError;
    
    protected Charset defaultCharset = Charset.forName("UTF-8");

    public Throwable getLastError() {
        return lastError;
    }

    public Cassandra.Client getClient() {
        return client;
    }

    public void setClient(Cassandra.Client client) {
        this.client = client;
    }

    public ConnectionManager getParentManager() {
        return parentManager;
    }

    public void setParentManager(ConnectionManager parentManager) {
        this.parentManager = parentManager;
    }

    @Override
    public void setKeyspace(String keyspace) {
        this.keyspace = keyspace;
        try {
            if (client != null) {
                client.set_keyspace(keyspace);
            }
        } catch (Exception e) {
            log.log(Level.SEVERE, "Cannot set new keyspace to use with connection.", e);
            throw new ConnectionException(e.getMessage(), e);
        }
    }

    @Override
    public String getKeyspace() {
        return keyspace;
    }

    @Override
    public void setLimit(int limit) {
        this.defaultLimit = limit;
    }

    @Override
    public int getLimit() {
        return defaultLimit;
    }

    @Override
    public void setConsistencyLevel(ConsistencyLevel level) {
        this.defaultLevel = level;
    }

    @Override
    public ConsistencyLevel getConsistencyLevel() {
        return defaultLevel;
    }

    @SuppressWarnings("unchecked")
	@Override
    public List<KeySlice> get(final String family, final Condition condition) {
        return (List<KeySlice>) tryExec(new Operation() {
            @Override
            public Object execute() throws Exception {
                return client.get_range_slices(new ColumnParent(family), condition.getPredicate(), condition.getRange(), defaultLevel);
            }
        });
    }

    @Override
    public Result query(final String query) {
//        CqlResult result = (CqlResult) tryExec(new Operation() {
//            @Override
//            public Object execute() throws Exception {
//                return client.execute_cql_query(ByteBuffer.wrap(query.getBytes(defaultCharset)), Compression.NONE);
//            }
//        });
        return null;
    }

    @Override
    public void insert(String columnFamily, ByteBuffer rowKey, List<Column> columns) {
        List<Mutation> mutationList = new ArrayList<Mutation>(columns.size());
        for (Column c : columns) {
            Mutation m = new Mutation();
            ColumnOrSuperColumn cs = new ColumnOrSuperColumn();
            cs.setColumn(c);
            m.setColumn_or_supercolumn(cs);
            mutationList.add(m);
        }
        prepareAndExecMutation(columnFamily, rowKey, mutationList);
    }

    @Override
    public void insert(String superColumnFamily, ByteBuffer rowKey, Map<ByteBuffer, List<Column>> columns) {
        List<Mutation> mutationList = new ArrayList<Mutation>(columns.size());
        for (ByteBuffer key : columns.keySet()) {
            Mutation m = new Mutation();
            ColumnOrSuperColumn data = new ColumnOrSuperColumn();
            SuperColumn sc = new SuperColumn();
            sc.setName(key);
            sc.setColumns(columns.get(key));
            data.setSuper_column(sc);
            m.setColumn_or_supercolumn(data);
            mutationList.add(m);
        }
        prepareAndExecMutation(superColumnFamily, rowKey, mutationList);
    }

    protected void prepareAndExecMutation(String columnFamily, ByteBuffer rowKey, List<Mutation> mutationList) {
        Map<String, List<Mutation>> dataMap = new HashMap<String, List<Mutation>>();
        dataMap.put(columnFamily, mutationList);
        final Map<ByteBuffer, Map<String, List<Mutation>>> mutationMap = new HashMap<ByteBuffer, Map<String, List<Mutation>>>();
        mutationMap.put(rowKey, dataMap);
        tryExec(new Operation() {
            @Override
            public String execute() throws TException, TimedOutException, InvalidRequestException, UnavailableException {
                client.batch_mutate(mutationMap, defaultLevel);
                return null;
            }
        });
    }
    
    @Override
    public void mutate(final Map<ByteBuffer, Map<String, List<Mutation>>> mutationMap) {
        tryExec(new Operation() {
            @Override
            public String execute() throws Exception {
                client.batch_mutate(mutationMap, defaultLevel);
                return null;
            }
        });
    }

    @Override
    public void finish() {
        //Input and output protocols seem to be the same in current implementation.
        parentManager.finish(client.getInputProtocol());
    }

    public void close() {
        client.getInputProtocol().getTransport().close();
    }
    
    protected Object tryExec(Operation op) {
        try {
            return op.execute();
        } catch (Exception e) {
            lastError = e;
            log.log(Level.SEVERE, "Error interacting with database.", e);
            throw new ConnectionException("Error interacting with database.", e);
        }
    }

    protected interface Operation {
        public Object execute() throws Exception ;
    }
}
